﻿using AsmShaderConverter.AsmShader;
using System;
using System.Collections.Generic;
using System.IO;
using System.Text;

namespace AsmShaderConverter.HlslShader
{
    public enum HlslShaderTy
    {
        PixelShader,
        VertexShader
    }

    public class HlslShaderFile
    {
        public HlslStruct InputStruct { get; set; }
        public HlslStruct OutputStruct { get; set; }

        private List<string> headerLines;
        private List<ConstantValue> constanRegisters;

        private List<IConvertToHlsl> Functions;

        public HlslShaderTy ShaderType { get; set; }

        public HlslShaderFile()
        {
            this.headerLines = new List<string>();
            this.constanRegisters = new List<ConstantValue>();
            this.Functions = new List<IConvertToHlsl>();
            this.InputStruct = new HlslStruct();
            this.OutputStruct = new HlslStruct();
        }

        public void AddHeadLines(string line)
        {
            this.headerLines.Add(line + "\n");
        }

        public void AddInput(AsmShaderInputOutput register)
        {
            string registertype = register.RegisterName.Substring(0, 1);
            switch (registertype)
            {
                case "v": // Input Register
                    this.InputStruct.AddRegister(register);
                    break;

                case "o": // Output Register
                    this.OutputStruct.AddRegister(register);
                    break;

                default:
                    throw new Exception("Unknown Register Type: " + register.RegisterName);
            }
        }

        internal void AddConstantValues(ConstantValue item)
        {
            this.constanRegisters.Add(item);
        }

        public void AddFunction(IConvertToHlsl function)
        {
            this.Functions.Add(function);
        }

        public string GetHeaderLines()
        {
            StringBuilder sb = new StringBuilder();

            foreach (var item in this.headerLines)
            {
                sb.AppendLine(item);
            }

            return sb.ToString();
        }

        public string GetInputLinesInput(AsmShaderFile asmShader, HlslShaderFile hlslVS)
        {
            StringBuilder sb = new StringBuilder();
            string mainMethod = string.Empty;
            string targetSemantic = string.Empty;

            if (this.InputStruct.StructTypeName == string.Empty)
                throw new Exception("structName must have a value");
            if (this.InputStruct.VariableName == string.Empty)
                throw new Exception("variableName must have a value");

            if (asmShader.ShaderType == AsmShaderHelper.ShaderType.PixelShader)
            {
                mainMethod = "ps_main";
                targetSemantic = ": SV_Target";
            }

            if (asmShader.ShaderType == AsmShaderHelper.ShaderType.VertexShader)
            {
                mainMethod = "vs_main";
                targetSemantic = "";
            }
            string InputStruct = CreateCodeForInputStruct(asmShader, hlslVS);
            sb.AppendLine(InputStruct);

            // Output structs
            if (asmShader.ShaderType == AsmShaderHelper.ShaderType.VertexShader) // Only add Output struct to VertexShader nor PixelShader
            {
                var outputStructs = CreateCodeForOutputStruct(asmShader);
                sb.AppendLine(outputStructs);
            }

            // Funktion header 'string NameOfFunction(string parameter)'
            sb.AppendLine(string.Format("{0} {1}({2} {3}) {4}\r\n{{\r\n", this.OutputStruct.StructTypeName, mainMethod, this.InputStruct.StructTypeName, this.InputStruct.VariableName, targetSemantic)); //

            return sb.ToString();
        }

        public string CreateCodeForInputStruct(AsmShaderFile asmShader, HlslShaderFile hlslVS)
        {
            StringBuilder sb2 = new StringBuilder();

            if (asmShader.ShaderType == AsmShaderHelper.ShaderType.PixelShader)
            {
                // Add missing registers to input struct for pixelshader
                if (hlslVS != null)
                {
                    AsmShaderConverter.AsmShader.AsmShaderInputOutput.AddMissingRegisterToPixelShaderInput(this.InputStruct.GetRegisters(), hlslVS.OutputStruct.GetRegisters());
                }
                else
                    throw new Exception("Vertex Shader should not be null. VS-Output struct must compare to Pixelshader Input Struct.");
            }

            sb2.AppendFormat("struct {0} \r\n{{\r\n", this.InputStruct.StructTypeName); // Start of struct

            var inputRegisters = this.InputStruct.GetRegisters();
            for (int i = 0; i < inputRegisters.Count; i++)
            {
                sb2.Append(inputRegisters[i].ToHLSLWitchComment(asmShader).Replace(this.InputStruct.VariableName + ".", ""));

                // if (i < this.ps_Inputs.Count)
                sb2.AppendLine(";");
            }
            sb2.AppendLine("};\n\n\n");
            string testcfdsf = sb2.ToString();
            return testcfdsf;
        }

        internal void SetInputStructName(string type, string name)
        {
            this.InputStruct.StructTypeName = type;
            this.InputStruct.VariableName = name;
        }

        internal void SetOutputStructName(string type, string name)
        {
            this.OutputStruct.StructTypeName = type;
            this.OutputStruct.VariableName = name;
        }

        public string GetEndOfShaderCode(AsmShaderFile asm)
        {
            StringBuilder sb = new StringBuilder();
            switch (asm.ShaderType)
            {
                case AsmShaderHelper.ShaderType.PixelShader:
                    sb.AppendLine("return " + GetReturnVariableForPixelShader(asm) + ";");
                    sb.AppendLine("}\n");
                    return sb.ToString();

                case AsmShaderHelper.ShaderType.VertexShader:
                    sb.AppendLine("return " + this.OutputStruct.VariableName + ";");
                    sb.AppendLine("}\n");
                    return sb.ToString();

                default:
                    throw new Exception("No Shadertype given to HlslShaderFile");
            }
        }

        internal string CreateShaderCode(AsmShaderFile asmShader)
        {
            StringBuilder sb = new StringBuilder();

            // Add all constantregisters.
            foreach (var item in this.constanRegisters)
            {
                sb.AppendLine(item.ToHLSLWitchComment(asmShader));
            }
            sb.AppendLine("\r\n");

            // Initializing line for output struct.
            sb.AppendLine(string.Format("{0} {1} = ({0})0;", this.OutputStruct.StructTypeName, this.OutputStruct.VariableName));

            // local registers.
            sb.AppendLine(CreateCodeForLocalRegister(asmShader));

            sb.AppendLine("\r\n");

            // Insert all functions
            foreach (var item in this.Functions)
            {
                sb.AppendLine(item.ToHLSLWitchComment(asmShader));
            }

            // Add 'return XXX }'-line to shadercode.
            sb.AppendLine(GetEndOfShaderCode(asmShader));

            return sb.ToString();
        }

        private string CreateCodeForLocalRegister(AsmShaderFile asm)
        {
            // Get All temp registered (r0-n)
            List<string> allLocalRegistered = new List<string>();
            StringBuilder sb2 = new StringBuilder();

            foreach (var item in this.Functions)
            {
                var register = item.GetDestinationRegister();
                if (register != null)
                    if (register.GetRegisterName(asm) == register.OriginalRegisterName)
                        if (register.OriginalRegisterName.StartsWith("r"))//|| register.OriginalRegisterName.StartsWith("o"))
                            if (!allLocalRegistered.Contains(register.OriginalRegisterName))
                                allLocalRegistered.Add(register.OriginalRegisterName);
            }
            allLocalRegistered.Sort();
            foreach (var item in allLocalRegistered)
            {
                sb2.AppendLine(string.Format("float4 {0} = (float4)0;", item));
            }
            return sb2.ToString();
        }

        private string GetReturnVariableForPixelShader(AsmShaderFile asm)
        {
            // Get All temp registered (r0-n)
            List<string> allLocalRegistered = new List<string>();
            StringBuilder sb2 = new StringBuilder();

            foreach (var item in this.Functions)
            {
                var register = item.GetDestinationRegister();
                if (register != null)
                    if (register.GetRegisterName(asm) == register.OriginalRegisterName)
                        if (register.OriginalRegisterName.StartsWith("o"))
                            if (!allLocalRegistered.Contains(register.OriginalRegisterName))
                                return register.OriginalRegisterName;
            }
            throw new Exception("No return Value For PixelShader found!");
        }

        public string CreateCodeForOutputStruct(AsmShaderFile asm)
        {
            StringBuilder sb2 = new StringBuilder();

            sb2.AppendLine(string.Format("struct {0}\r\n{{", this.OutputStruct.StructTypeName));

            foreach (var item in this.OutputStruct.GetRegisters())
            {
                string text = item.ToHLSLWitchComment(asm).Replace(this.OutputStruct.VariableName + ".", "");
                sb2.AppendLine(text + ";");
            }

            sb2.AppendLine("};\n");

            return sb2.ToString();
        }

        internal void SaveToFile(string filename, HlslShaderFile hlslVS)
        {
            StreamWriter sw = new StreamWriter(filename);

            sw.Write(GetHeaderLines());
            sw.Write(GetInputLinesInput(this.AsmShader, hlslVS));
            sw.Write(CreateShaderCode(this.AsmShader));
            // is already in GetShaderCpde : sw.Write(GetEndOfShaderCode(this.AsmShader));

            sw.Close();
        }

        public AsmShaderFile AsmShader { get; set; }

        internal string GetInputLinesInput(AsmShaderFile psAsm)
        {
            return GetInputLinesInput(psAsm, null);
        }

        internal List<HlslShader.Sampler.HlslSampler> GetUsedSampler()
        {
            //{
            //    Dictionary<string, HlslShader.Sampler.HlslSampler> samplers = new Dictionary<string, Sampler.HlslSampler>();

            return this.AsmShader.GetSamplerRegister();
            //{
            //}
            //return samplers;
        }
    }
}